{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Solving constrained optimization problems\n",
    "\n",
    "---\n",
    "\n",
    "In this example, we consider the following constrained optimization problem:\n",
    "\n",
    "$$\n",
    "\\begin{array}{r c}\n",
    "\\text{maximize }   & a + b + c \\\\\n",
    "\\text{subject to } & 2a + 3b \\leq 45 \\\\\n",
    "                   & 5a + 2c \\leq 75 \\\\\n",
    "                   & 3b + c \\leq 50 \\\\\n",
    "                   & -100 \\leq a \\leq 100 \\\\\n",
    "                   & -100 \\leq b \\leq 100 \\\\\n",
    "                   & -100 \\leq c \\leq 100 \\\\\n",
    "\\end{array}\n",
    "$$\n",
    "\n",
    "We will now solve this optimization problem using:\n",
    "\n",
    "- Evolutionary computation (using PGPE with constraint penalties)\n",
    "- Penalty method (using gradient-based search with ClipUp, with penalty multipliers iteratively incremented)\n",
    "- Interior points method (with the help of the functional API of PyTorch and the log-barrier function of EvoTorch)\n",
    "\n",
    "This notebook demonstrates:\n",
    "\n",
    "- How to use functional algorithms of EvoTorch (evolutionary and gradient-based)\n",
    "- How to use constraint penalization utilities of EvoTorch and use them as building blocks for constrained optimization, both for evolutionary and for gradient-based optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from evotorch.algorithms.functional import pgpe, pgpe_ask, pgpe_tell, clipup, clipup_ask, clipup_tell\n",
    "from evotorch.decorators import expects_ndim\n",
    "from evotorch.tools import penalty, log_barrier\n",
    "\n",
    "import torch\n",
    "import torch.func as tfunc\n",
    "from typing import Union\n",
    "from functools import partial\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Fitness function implementation\n",
    "\n",
    "We begin by implementing our fitness function.\n",
    "This fitness function has two arguments:\n",
    "\n",
    "- `penalty_multiplier`: Expected as a scalar, this value represents the multiplier for the negative penalty quantities that will be added onto the fitness. Higher values for `penalty_multiplier` will result in stronger penalizations.\n",
    "- `x`: A 1-dimensional tensor, that represents the solution that will be evaluated.\n",
    "\n",
    "Notice how the fitness function below is decorated by `@expects_ndim(0, 1)`. This decorators declares that the first positional argument (`penalty_multipliers`) expects a 0-dimensional tensor, and the second positional argument (`x`) expects a 1-dimensional tensor. If any of these arguments are given with more dimensions, those extra dimensions will be considered as batch dimensions by the decorated function. This auto-batching feature is very helpful, because it allows the decorated function `f` to work with not just a single solution, but with a population of solutions (or a batch of populations of solutions) when `x` is given as an n-dimensional tensor with `n>1`. We will use this auto-batching feature with PGPE."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@expects_ndim(0, 1)\n",
    "def f(penalty_multiplier: torch.Tensor, x: torch.Tensor) -> torch.Tensor:\n",
    "    a = x[0]\n",
    "    b = x[1]\n",
    "    c = x[2]\n",
    "    objective = a + b + c\n",
    "\n",
    "    constraints = [\n",
    "        [(2 * a) + (3 * b), \"<=\", 45],\n",
    "        [(5 * a) + (2 * c), \"<=\", 75],\n",
    "        [(3 * b) + c, \"<=\", 50],\n",
    "        [a, \">=\", -100],\n",
    "        [a, \"<=\", 100],\n",
    "        [b, \">=\", -100],\n",
    "        [b, \"<=\", 100],\n",
    "        [c, \">=\", -100],\n",
    "        [c, \"<=\", 100],\n",
    "    ]\n",
    "\n",
    "    penalty_amount = 0.0\n",
    "    for lhs, op, rhs in constraints:\n",
    "        # For each constraint, we add a penalty (if there is violation)\n",
    "        penalty_amount = penalty_amount + penalty(\n",
    "            # Left-hand-side, comparison operator, and the right-hand-side:\n",
    "            lhs,\n",
    "            op,\n",
    "            rhs,\n",
    "            #\n",
    "            # Because this is a function we wish to maximize, the penalties should be in the opposite direction.\n",
    "            # Therefore, we declare the sign of the penalties as \"-\", making them negative quantities:\n",
    "            penalty_sign=\"-\",\n",
    "            #\n",
    "            # There will be a penalty in the form: (linear * amount_of_violation)\n",
    "            linear=1.0,\n",
    "            #\n",
    "            # There will also be a penalty in the form: (amount_of_violation ** exp)\n",
    "            exp=2.0,\n",
    "            #\n",
    "            # The exponential penalties are not allowed to exceed this amount:\n",
    "            exp_inf=5000.0,\n",
    "            #\n",
    "            # There will also be a penalty in the form: (step if amount_of_violation > 0 else 0)\n",
    "            step=10000.0,\n",
    "        )\n",
    "\n",
    "    # Scale the accumulated penalty by `penalty_multiplier`, add it onto the objective, then return it\n",
    "    return objective + (penalty_amount * penalty_multiplier)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Evolutionary computation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "As the evolutionary algorithm, we use the functional `pgpe`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "pgpe_state = pgpe(\n",
    "    # We want to maximize the evaluation results:\n",
    "    objective_sense=\"max\",\n",
    "\n",
    "    # The center point of the initial search distribution is given as 0 vector:\n",
    "    center_init=torch.zeros(3),\n",
    "\n",
    "    # Learning rates for the center and the standard deviation of the\n",
    "    # search distribution:\n",
    "    center_learning_rate=0.1,\n",
    "    stdev_learning_rate=0.1,\n",
    "\n",
    "    # The ranking method is \"centered\", which ranks the worst solution as\n",
    "    # -0.5, and the best solution as +0.5:\n",
    "    ranking_method=\"centered\",\n",
    "\n",
    "    # We use the ClipUp optimizer:\n",
    "    optimizer=\"clipup\",\n",
    "\n",
    "    # In the case of this example problem, we use the following max_speed:\n",
    "    optimizer_config={\"max_speed\": 1.0},\n",
    "\n",
    "    # The standard deviation for the initial search distribution:\n",
    "    stdev_init=100.0,\n",
    "\n",
    "    # Relative maximum allowed change for the standard deviation.\n",
    "    # 0.2 means that a standard deviation value cannot change more than 20%\n",
    "    # of its original value.\n",
    "    stdev_max_change=0.2,\n",
    "\n",
    "    # Minimum standard deviation value. The standard deviation is cannot shrink\n",
    "    # to values lower than this:\n",
    "    stdev_min=0.01,\n",
    ")\n",
    "\n",
    "pgpe_state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "Main loop of the evolutionary computation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_generations = 1000  # For how many iterations will PGPE work\n",
    "\n",
    "time_interval = 1  # With this period, we will print the status\n",
    "last_report_time = datetime.now()\n",
    "\n",
    "# Initialize the variables that will store the best solution and the best evaluation result\n",
    "pgpe_best_eval = float(\"-inf\")\n",
    "pgpe_best = None\n",
    "\n",
    "for generation in range(1, 1 + num_generations):\n",
    "    # Ask for a population from PGPE\n",
    "    population = pgpe_ask(pgpe_state, popsize=1000)\n",
    "\n",
    "    # Compute the fitnesses\n",
    "    fitnesses = f(1, population)\n",
    "\n",
    "    # Inform PGPE of the latest fitnesses, and get its next state\n",
    "    pgpe_state = pgpe_tell(pgpe_state, population, fitnesses)\n",
    "\n",
    "    # From the most recent population and fitnesses, update the best solution and the best evaluation result\n",
    "    pop_best_index = torch.argmax(fitnesses)\n",
    "    pop_best = population[pop_best_index]\n",
    "    pop_best_eval = fitnesses[pop_best_index]\n",
    "    if pop_best_eval > pgpe_best_eval:\n",
    "        pgpe_best_eval = pop_best_eval\n",
    "        pgpe_best = population[pop_best_index, :]\n",
    "\n",
    "    # If it is time to report, print the status\n",
    "    tnow = datetime.now()\n",
    "    if ((tnow - last_report_time).total_seconds() > time_interval) or (generation == num_generations):\n",
    "        print(\"best solution:\", pgpe_best, \"best evaluation:\", pgpe_best_eval)\n",
    "        last_report_time = tnow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "Best solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "pgpe_best, pgpe_best_eval"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Penalty method\n",
    "\n",
    "As an alternative to evolutionary computation, we now use the gradient-based penalty method below.\n",
    "The main points of the penalty method are as follows:\n",
    "\n",
    "- Use a gradient-based search on a fitness function augmented by penalties (i.e. by penalty terms that are computed according to how much the constraints are penalized)\n",
    "- Periodically increase the multipliers for the penalty terms (when the penalty multiplier is increased, we use the previous optimization's solution as the new starting point)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Try with these penalty multipliers, ordered from small to large:\n",
    "penalty_multipliers = [0.1, 1, 4]\n",
    "\n",
    "# For each penalty multiplier, we do the search for this number of iterations:\n",
    "num_iterations = 1000\n",
    "\n",
    "# Initialize the variables that will store the best solution and the best evaluation result\n",
    "clipup_best_eval = float(\"-inf\")\n",
    "clipup_best = None\n",
    "\n",
    "# Initialize the search point as a 0 vector:\n",
    "x = torch.zeros(3)\n",
    "\n",
    "for penalty_multiplier in penalty_multipliers:\n",
    "    # Initialize the ClipUp algorithm for the currently considered penalty multiplier\n",
    "    clipup_state = clipup(center_init=x, center_learning_rate=0.1, max_speed=1.0)\n",
    "\n",
    "    # Optimization loop for the current penalty multiplier\n",
    "    for iteration in range(1, 1 + num_iterations):\n",
    "        # Ask ClipUp for the current search point\n",
    "        x = clipup_ask(clipup_state)\n",
    "\n",
    "        # Compute the gradient and the fitness\n",
    "        gradient, fitness = tfunc.grad_and_value(partial(f, penalty_multiplier))(x)\n",
    "\n",
    "        # Update the best-known solution so far\n",
    "        if fitness > clipup_best_eval:\n",
    "            clipup_best_eval = fitness\n",
    "            clipup_best = x\n",
    "\n",
    "        # Inform ClipUp of the latest gradient, and get its next state\n",
    "        clipup_state = clipup_tell(clipup_state, follow_grad=gradient)\n",
    "\n",
    "    # After each optimization loop, print the best known solution\n",
    "    print(\"best solution:\", clipup_best, \" best evaluation:\", clipup_best_eval)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Interior points method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "Although, as its name implies, the main focus of EvoTorch is evolutionary computation, we now demonstrate that it is possible to implement an interior points method by combining:\n",
    "\n",
    "- `torch.func.grad`\n",
    "- `torch.func.hessian`\n",
    "- `evotorch.tools.log_barrier`\n",
    "\n",
    "The `grad(...)` and `hessian(...)` functions are basic building blocks for implementing a Newton-Raphson search. We penalize proximities to the infeasible regions with the help of `log_barrier(...)`.\n",
    "\n",
    "We begin with a modified implementation of our fitness function, where `evotorch.tools.penalty(...)` are replaced by `evotorch.tools.log_barrier(...)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "@expects_ndim(0, 1)\n",
    "def f_with_log_barriers(sharpness: torch.Tensor, x: torch.Tensor) -> torch.Tensor:\n",
    "    a = x[0]\n",
    "    b = x[1]\n",
    "    c = x[2]\n",
    "    objective = a + b + c\n",
    "\n",
    "    constraints = [\n",
    "        [(2 * a) + (3 * b), \"<=\", 45],\n",
    "        [(5 * a) + (2 * c), \"<=\", 75],\n",
    "        [(3 * b) + c, \"<=\", 50],\n",
    "        [a, \">=\", -100],\n",
    "        [a, \"<=\", 100],\n",
    "        [b, \">=\", -100],\n",
    "        [b, \"<=\", 100],\n",
    "        [c, \">=\", -100],\n",
    "        [c, \"<=\", 100],\n",
    "    ]\n",
    "\n",
    "    penalty_amount = 0.0\n",
    "    for lhs, op, rhs in constraints:\n",
    "        penalty_amount = penalty_amount + log_barrier(lhs, op, rhs, penalty_sign=\"-\", sharpness=sharpness)\n",
    "\n",
    "    # Return the objective with the log-barrier penalties added onto it.\n",
    "    # Notice that in the end, we are inverting the sign of the returned quantity.\n",
    "    # This will allow us to implement the Newton-Raphson's search method from a minimization perspective.\n",
    "    return -(objective + penalty_amount)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the search from the 0 vector:\n",
    "x = torch.zeros(3)\n",
    "\n",
    "# Try with these sharpness values for the log barrier, ordered from small to large:\n",
    "sharpness_values = [1, 10, 100, 1000, 10000]\n",
    "\n",
    "# Learning rate for when taking a step\n",
    "learning_rate = 0.1\n",
    "\n",
    "# An identity matrix multiplied by the constant below will be added to the Hessian matrix.\n",
    "# When a diagonal element of the Hessian matrix is 0 because of numerical issues, this trick will allow the\n",
    "# algorithm to still take a step.\n",
    "I_multiplier = 0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# With this interval (in seconds), we wish to report the status\n",
    "reporting_interval = 1\n",
    "last_report_time = datetime.now()\n",
    "\n",
    "# The interior points method will run for this amount of iterations:\n",
    "num_iterations = 200\n",
    "\n",
    "for sharpness in sharpness_values:\n",
    "    print(\"sharpness:\", sharpness)\n",
    "    print()\n",
    "\n",
    "    for iteration in range(1, 1 + num_iterations):\n",
    "        # Compute the gradient and the solution cost\n",
    "        g, cost = tfunc.grad_and_value(partial(f_with_log_barriers, sharpness))(x)\n",
    "\n",
    "        # Compute the Hessian matrix\n",
    "        H = tfunc.hessian(partial(f_with_log_barriers, sharpness))(x)\n",
    "\n",
    "        # Add the identity matrix multiplied by a constant to the Hessian\n",
    "        H = H + (I_multiplier * torch.eye(H.shape[-1]))\n",
    "\n",
    "        # Take the inverse of the Hessian matrix\n",
    "        invH = torch.linalg.inv(H)\n",
    "\n",
    "        # Move the center point\n",
    "        x = x - learning_rate * (invH @ g)\n",
    "\n",
    "        # If it is time to report, print the status\n",
    "        tnow = datetime.now()\n",
    "        if (tnow - last_report_time).total_seconds() > reporting_interval:\n",
    "            print(\"Iteration:\", iteration, \"  Gradient norm:\", torch.linalg.norm(g), \"  Solution cost:\", cost)\n",
    "            last_report_time = tnow\n",
    "\n",
    "    # Print the current search point\n",
    "    print()\n",
    "    print(\"x:\", x)\n",
    "    print()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
